# -*- coding: utf-8 -*-
"""
绘制：x=total_size，y=该 total_size 下 comm_time(us) 的平均值
- 在 __main__ 中定义 csv_path 与 comm_type
- 仅提取 comm_type 指定的行
- 生成折线图并保存为 PNG
"""

import os
import sys
import pandas as pd
import matplotlib.pyplot as plt
from typing import Optional, Union


REQUIRED_COLUMNS = [
    "comm_id","comm_size","comm_type","rank","root","time_stamp",
    "comm_time(us)","sendsize","sendcount","recvsize","recvcount",
    "src","dst","total_size","appearance_time"
]


def validate_columns(df: pd.DataFrame):
    missing = [c for c in REQUIRED_COLUMNS if c not in df.columns]
    if missing:
        raise ValueError(f"缺少必要列: {missing}\n实际列: {list(df.columns)}")


def load_and_filter(csv_path: str, comm_type: Union[str, int]) -> pd.DataFrame:
    # 读取
    try:
        df = pd.read_csv(csv_path)
    except UnicodeDecodeError:
        # 备用编码
        df = pd.read_csv(csv_path, encoding="utf-8-sig")
    validate_columns(df)

    # 过滤 comm_type（兼容数字/字符串，如 51 与 "51"）
    comm_type_norm = str(comm_type).strip()
    df["_comm_type_norm"] = df["comm_type"].astype(str).str.strip()
    filtered = df[df["_comm_type_norm"] == comm_type_norm].copy()
    df.drop(columns=["_comm_type_norm"], inplace=True, errors="ignore")
    if filtered.empty:
        raise ValueError(f"在文件中未找到 comm_type == '{comm_type}' 的数据行。")

    # 类型转换与清洗
    # total_size 可能是整数/浮点，comm_time(us) 应为数值
    for col in ["total_size", "comm_time(us)"]:
        filtered[col] = pd.to_numeric(filtered[col], errors="coerce")
    filtered = filtered.dropna(subset=["total_size", "comm_time(us)"])

    return filtered


def aggregate_mean_by_total_size(df: pd.DataFrame) -> pd.DataFrame:
    agg = (
        df.groupby("total_size", as_index=False)["comm_time(us)"]
        .mean()
        .rename(columns={"comm_time(us)": "mean_comm_time_us"})
        .sort_values("total_size")
    )
    if agg.empty:
        raise ValueError("聚合结果为空，检查 total_size 与 comm_time(us) 是否为有效数值。")
    return agg


def plot_and_save(agg_df: pd.DataFrame, csv_path: str, comm_type: Union[str, int], out_dir: Optional[str] = None) -> str:
    plt.figure(figsize=(8, 5))
    plt.plot(
        agg_df["total_size"],
        agg_df["mean_comm_time_us"],
        marker="o",
        linestyle="-",
        color="#1f77b4",
        label=f"comm_type={comm_type}"
    )
    plt.xlabel("total_size")
    plt.ylabel("comm_time(us)")
    plt.title("total_size vs avg_comm_time(us)")
    plt.grid(True, linestyle="--", alpha=0.4)
    plt.legend()

    # 选择保存目录：优先使用传入的 out_dir，否则使用 csv 所在目录
    save_dir = out_dir or (os.path.dirname(csv_path) or ".")
    os.makedirs(save_dir, exist_ok=True)
    base = os.path.splitext(os.path.basename(csv_path))[0]
    out_file = os.path.join(save_dir, f"{base}_commtype_{comm_type}_mean_comm_time_vs_total_size.png")
    plt.tight_layout()
    plt.savefig(out_file, dpi=150)
    plt.close()
    return out_file


def main():
    # 1) 在 __main__ 中定义 csv_path（你的CSV路径）
    csv_path = r"F:\PostGraduate\Point-to-Point-DATA\DIFF_ATOM\2node\2node-32proc-10interation-2000000atom-20250925_155735\log-0_processed.csv"

    # 3) 在 __main__ 中定义 comm_type（只提取相应的数据行）
    # 示例：根据你的数据实际值设置，如 "MPI_Allreduce"、"MPI_Bcast"、"SendRecv" 等
    comm_type = 51

    # 可选输出目录；设为 None 或空字符串则默认保存到 csv 同目录
    output_dir = r"F:\PostGraduate\Point-to-Point-Code\App_Prediction\code\tools_output\point_size_and_time_relation_fron_one_csv"

    if not os.path.exists(csv_path):
        print(f"文件不存在：{csv_path}", file=sys.stderr)
        sys.exit(1)

    try:
        df = load_and_filter(csv_path, comm_type)
        agg_df = aggregate_mean_by_total_size(df)
        out_file = plot_and_save(agg_df, csv_path, comm_type, out_dir=output_dir)
    except Exception as e:
        print(f"处理失败：{e}", file=sys.stderr)
        sys.exit(2)

    # 简要输出统计信息
    print(f"已生成图像：{out_file}")
    print(f"数据点数量：{len(agg_df)}")
    print("前5行汇总：")
    print(agg_df.head().to_string(index=False))


if __name__ == "__main__":
    main()
